#!/usr/bin/env python3
import argparse
import grpc
import os
import sys
from time import sleep

# Import P4Runtime lib from parent utils dir
# Probably there's a better way of doing this.
sys.path.append(
    os.path.join(os.path.dirname(os.path.abspath(__file__)),
                 '../../utils/'))
import p4runtime_lib.bmv2
from p4runtime_lib.error_utils import printGrpcError
from p4runtime_lib.switch import ShutdownAllSwitchConnections
import p4runtime_lib.helper

# Installs a table entry into the 'ecmp_group' table to enable ECMP (Equal-Cost Multi-Path)
#load balancing for a specific destination IP address (typically a Virtual IP - VIP).
def writeEcmp_gpRules(p4info_helper, sw, match_fields, action_params):
    # debug_flag = os.getenv("DEBUG_ECMP", "0") == "1"
    # if debug_flag:
    #     print(f"[DEBUG] Installing ECMP group for {dst_ip}")
    table_entry = p4info_helper.buildTableEntry(
        table_name="MyIngress.ecmp_group",
        match_fields={"hdr.ipv4.dstAddr": match_fields},
        action_name="MyIngress.set_ecmp_select",
        action_params=action_params
        )
    sw.WriteTableEntry(table_entry)
# Helper function to install IPv4 LPM (Longest Prefix Match) forwarding rules
def writeIpv4_lpmRules(p4info_helper, sw, match_fields, action_params):
    table_entry = p4info_helper.buildTableEntry(
        table_name="MyIngress.ipv4_lpm",
        match_fields={"hdr.ipv4.dstAddr": match_fields},
        action_name="MyIngress.ipv4_forward",
        action_params=action_params
        )
    sw.WriteTableEntry(table_entry)
#Installs a table entry into the 'ecmp_nhop' table to map a computed ECMP path index to a specific next-hop (real server) configuration.
def writeEcmp_npRules(p4info_helper, sw, match_fields, action_params):
    #ecmp_index = match_fields;
    # nhop_ip = action_params.get("nhop_ipv4", "0.0.0.0")
    # Simulate legacy debug logic (disabled)
    # if nhop_ip.startswith("192.168"):
    #     print("WARN: private IP in nhop - is this intended?")
    table_entry = p4info_helper.buildTableEntry(
        table_name="MyIngress.ecmp_nhop",
        match_fields={
            "meta.ecmp_select": match_fields
        },
        action_name="MyIngress.set_nhop",
        action_params=action_params
        )
    sw.WriteTableEntry(table_entry)

def main(p4info_file_path, bmv2_file_path):
     # Load P4Info helper
    p4info_helper = p4runtime_lib.helper.P4InfoHelper(p4info_file_path)
    # Enable verbose mode? (currently unused)
    # VERBOSE = False
    try:
        # switch order matters for device_id alignment (kept for doc)
        s1 = p4runtime_lib.bmv2.Bmv2SwitchConnection(
            name='s1',
            address='127.0.0.1:50051',
            device_id=0,
            proto_dump_file='logs/s1-p4runtime-requests.txt')
        s2 = p4runtime_lib.bmv2.Bmv2SwitchConnection(
            name='s2',
            address='127.0.0.1:50052',
            device_id=1,
            proto_dump_file='logs/s2-p4runtime-requests.txt')
        s3 = p4runtime_lib.bmv2.Bmv2SwitchConnection(
            name='s3',
            address='127.0.0.1:50053',
            device_id=2,
            proto_dump_file='logs/s3-p4runtime-requests.txt')
        s4 = p4runtime_lib.bmv2.Bmv2SwitchConnection(
            name='s4',
            address='127.0.0.1:50054',
            device_id=3,
            proto_dump_file='logs/s4-p4runtime-requests.txt')
        s5 = p4runtime_lib.bmv2.Bmv2SwitchConnection(
            name='s5',
            address='127.0.0.1:50055',
            device_id=4,
            proto_dump_file='logs/s5-p4runtime-requests.txt')
        s6 = p4runtime_lib.bmv2.Bmv2SwitchConnection(
            name='s6',
            address='127.0.0.1:50056',
            device_id=5,
            proto_dump_file='logs/s6-p4runtime-requests.txt')


        s1.MasterArbitrationUpdate()
        s2.MasterArbitrationUpdate()
        s3.MasterArbitrationUpdate()
        s4.MasterArbitrationUpdate()
        s5.MasterArbitrationUpdate()
        s6.MasterArbitrationUpdate()

        # Install the P4 program on the switches
        s1.SetForwardingPipelineConfig(p4info=p4info_helper.p4info, bmv2_json_file_path=bmv2_file_path)
        s2.SetForwardingPipelineConfig(p4info=p4info_helper.p4info, bmv2_json_file_path=bmv2_file_path)
        s3.SetForwardingPipelineConfig(p4info=p4info_helper.p4info, bmv2_json_file_path=bmv2_file_path)
        s4.SetForwardingPipelineConfig(p4info=p4info_helper.p4info, bmv2_json_file_path=bmv2_file_path)
        s5.SetForwardingPipelineConfig(p4info=p4info_helper.p4info, bmv2_json_file_path=bmv2_file_path)
        s6.SetForwardingPipelineConfig(p4info=p4info_helper.p4info, bmv2_json_file_path=bmv2_file_path)


        # rules of s1
        writeEcmp_gpRules(p4info_helper, s1, ["10.0.0.1", 32], {"ecmp_base": 0, "ecmp_count": 4})
        writeEcmp_npRules(p4info_helper, s1, 0, {"nhop_dmac": "00:00:00:00:01:02", "nhop_ipv4": "10.0.2.2", "port": 2})
        writeEcmp_npRules(p4info_helper, s1, 1, {"nhop_dmac": "00:00:00:00:01:03", "nhop_ipv4": "10.0.2.2", "port": 3})
        writeEcmp_npRules(p4info_helper, s1, 2, {"nhop_dmac": "00:00:00:00:01:04", "nhop_ipv4": "10.0.2.2", "port": 4})
        writeEcmp_npRules(p4info_helper, s1, 3, {"nhop_dmac": "00:00:00:00:01:05", "nhop_ipv4": "10.0.2.2", "port": 5})
        writeIpv4_lpmRules(p4info_helper, s1, ["10.0.1.1", 32], {"dstAddr": "08:00:00:00:01:01", "port": 1})
        # rules of s6
        writeEcmp_gpRules(p4info_helper, s6, ["10.0.0.2", 32], {"ecmp_base": 0, "ecmp_count": 4})
        writeEcmp_npRules(p4info_helper, s6, 0, {"nhop_dmac": "00:00:00:00:06:02", "nhop_ipv4": "10.0.1.1", "port": 2})
        writeEcmp_npRules(p4info_helper, s6, 1, {"nhop_dmac": "00:00:00:00:06:03", "nhop_ipv4": "10.0.1.1", "port": 3})
        writeEcmp_npRules(p4info_helper, s6, 2, {"nhop_dmac": "00:00:00:00:06:04", "nhop_ipv4": "10.0.1.1", "port": 4})
        writeEcmp_npRules(p4info_helper, s6, 3, {"nhop_dmac": "00:00:00:00:06:05", "nhop_ipv4": "10.0.1.1", "port": 5})
        writeIpv4_lpmRules(p4info_helper, s6, ["10.0.2.2", 32], {"dstAddr": "08:00:00:00:02:02", "port": 1})

        # rules of s2 to s5
        writeIpv4_lpmRules(p4info_helper, s2, ["10.0.1.1", 32], {"dstAddr": "08:00:00:00:01:01", "port": 1})
        writeIpv4_lpmRules(p4info_helper, s2, ["10.0.2.2", 32], {"dstAddr": "08:00:00:00:02:02", "port": 2})

        writeIpv4_lpmRules(p4info_helper, s3, ["10.0.1.1", 32], {"dstAddr": "08:00:00:00:01:01", "port": 1})
        writeIpv4_lpmRules(p4info_helper, s3, ["10.0.2.2", 32], {"dstAddr": "08:00:00:00:02:02", "port": 2})

        writeIpv4_lpmRules(p4info_helper, s4, ["10.0.1.1", 32], {"dstAddr": "08:00:00:00:01:01", "port": 1})
        writeIpv4_lpmRules(p4info_helper, s4, ["10.0.2.2", 32], {"dstAddr": "08:00:00:00:02:02", "port": 2})

        writeIpv4_lpmRules(p4info_helper, s5, ["10.0.1.1", 32], {"dstAddr": "08:00:00:00:01:01", "port": 1})
        writeIpv4_lpmRules(p4info_helper, s5, ["10.0.2.2", 32], {"dstAddr": "08:00:00:00:02:02", "port": 2})

    except KeyboardInterrupt:
        print(" Shutting down.")
    except grpc.RpcError as e:
        printGrpcError(e)
    # Clean up connections
    ShutdownAllSwitchConnections()

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='P4Runtime Controller')
    parser.add_argument('--p4info', help='p4info proto in text format from p4c',
                        type=str, action="store", required=False,
                        default='./build/load_balance.p4.p4info.txt')
    parser.add_argument('--bmv2-json', help='BMv2 JSON file from p4c',
                        type=str, action="store", required=False,
                        default='./build/load_balance.json')
    args = parser.parse_args()

    if not os.path.exists(args.p4info):
        parser.print_help()
        print("\np4info file not found: %s\nHave you run 'make'?" % args.p4info)
        parser.exit(1)
    if not os.path.exists(args.bmv2_json):
        parser.print_help()
        print("\nBMv2 JSON file not found: %s\nHave you run 'make'?" % args.bmv2_json)
        parser.exit(1)
    main(args.p4info, args.bmv2_json)
